//===----------------------------------------------------------------------===//
//
// Copyright (C) 2022 Sophgo Technologies Inc.  All rights reserved.
//
// SOPHON-STREAM is licensed under the 2-Clause BSD License except for the
// third-party components.
//
//===----------------------------------------------------------------------===//

#pragma once
#include <atomic>
#include <deque>
#include <limits>
#include <mutex>
#include <condition_variable>
#include <optional>
#include <utility>
#include <thread>
#include <vector>

#define PIPE_INFO 0     // 是否打印管道已满的log

#if PIPE_INFO
#include <iostream>
#endif

enum class DiscardMode {
  POP_FRONT,
  POP_BACK,
  BLOCK,
};

template <typename T>
class DataPipe{
public:
  DataPipe() noexcept {}
  DataPipe(const char* name, int max_len): name_(name), max_size_(static_cast<size_t>(max_len)) {}
  DataPipe(const char* name, int max_len, DiscardMode mode): name_(name), max_size_(static_cast<size_t>(max_len)), discard_mode_(mode) {}

  DataPipe(DataPipe&& other) noexcept
    : data_que_(std::move(other.data_que_)),
      name_(other.name_),
      max_size_(other.max_size_.load()),
      discard_mode_(other.discard_mode_) {
    // mutex 和 condition_variable 不能移动，只能新建空的
  }

  // Move assignment operator
  DataPipe& operator=(DataPipe&& other) noexcept {
    if (this != &other) {
      std::lock_guard<std::mutex> lock(data_que_mtx_);
      data_que_ = std::move(other.data_que_);
      name_ = other.name_;
      max_size_ = other.max_size_.load();
      discard_mode_ = other.discard_mode_;
      // mutex 和 cv 保持各自的新对象
    }
    return *this;
  }

  std::size_t size() const {
    std::lock_guard<std::mutex> lock(data_que_mtx_);
    return data_que_.size();
  }

  template <typename U>
  void emplace_back(U&& data){
    while(true){
      std::unique_lock<std::mutex> lock(data_que_mtx_);
      if(data_que_.size() >= max_size_) {
#if PIPE_INFO
        std::cout << name_ << " pipe full" << std::endl;
#endif
        if (discard_mode_ == DiscardMode::POP_FRONT) {
          data_que_.pop_front();
        }
        else if (discard_mode_ == DiscardMode::POP_BACK) {
          data_que_.pop_back();
        }
        else {
          que_full_cv_.wait(lock, [this]() {return this->data_que_.size() < max_size_;});
        }
        continue;
      }
      else {
        data_que_.emplace_back(std::forward<U>(data));
        que_empty_cv_.notify_one();
        break;
      }
    }
  }

  template <typename U>
  bool try_emplace_back(U&& data) {
    std::unique_lock<std::mutex> lk(data_que_mtx_);
    if (data_que_.size() >= max_size_) {
      return false;
    }

    data_que_.emplace_back(std::forward<U>(data));
    que_empty_cv_.notify_one();
    return true;
  }

  auto block_pop_front() -> T {
    while(true){
      std::unique_lock<std::mutex> lock(data_que_mtx_);
      if(data_que_.size() == 0) {
        que_empty_cv_.wait(lock, [this]() {return !this->data_que_.empty();});
        continue;
      }
      else {
        auto ret = std::move(data_que_.front());
        data_que_.pop_front();
        que_full_cv_.notify_one();
        return ret;
      }
    }
  }

  auto block_pop_front_for(std::chrono::duration<int, std::milli> ms) -> std::optional<T> {
    std::unique_lock<std::mutex> lk(data_que_mtx_);
    if (data_que_.empty()) {
      bool ready = que_empty_cv_.wait_for(lk, ms, [this]() {
        return !data_que_.empty();
      });
      if (!ready) {
        return std::nullopt;
      }
    }
    auto ret = std::make_optional(std::move(data_que_.front()));
    data_que_.pop_front();
    que_full_cv_.notify_one();
    return ret;
  }

  // block wait unless flag evals true
  auto block_pop_front_unless(std::atomic_bool& flag, std::memory_order order = std::memory_order_acquire) -> std::optional<T> {
    std::unique_lock<std::mutex> lk(data_que_mtx_);
    if (data_que_.empty()) {
      que_empty_cv_.wait(lk, [this, &flag, order]() {
        return !data_que_.empty() || flag.load(order);
      });
      if (data_que_.empty()) {
        que_empty_cv_.notify_one();
        return std::nullopt;
      }
    }
    auto ret = std::make_optional(std::move(data_que_.front()));
    data_que_.pop_front();
    que_full_cv_.notify_one();
    return ret;
  }

  auto try_pop_front() -> std::optional<T> {
    for (;;) {
      // fast path: try lock-free check
      if (data_que_mtx_.try_lock()) {
        if (!data_que_.empty()) {
          auto ret = std::make_optional(std::move(data_que_.front()));
          data_que_.pop_front();
          data_que_mtx_.unlock();
          que_full_cv_.notify_one();
          return ret;
        } else {
          data_que_mtx_.unlock();
          return std::nullopt;
        }
      }

      // failed to acquire lock: backoff
      static thread_local int spin = 1;
      for (int i = 0; i < spin; i++) {
        // hint to CPU that we are spinning
        std::this_thread::yield();
      }
      // exponential backoff, capped
      spin = std::min(spin * 2, 1 << 10);
    }
  }

  // auto try_pop_front() -> std::optional<T> {
  //   std::unique_lock<std::mutex> lk(data_que_mtx_);
  //   if (data_que_.empty()) {
  //     return std::nullopt;
  //   } else {
  //     auto ret = std::make_optional(std::move(data_que_.front()));
  //     data_que_.pop_front();
  //     que_full_cv_.notify_one();
  //     return ret;
  //   }
  // }

  size_t pop_front_all(std::vector<T>& buffer) {
    std::unique_lock<std::mutex> lk(data_que_mtx_);
    buffer.reserve(buffer.size() + data_que_.size());
    size_t cnt = 0;
    while (!data_que_.empty()) {
      buffer.emplace_back(std::move(data_que_.front()));
      data_que_.pop_front();
      cnt++;
    }
    return cnt;
  }

  void set_mode(DiscardMode mode) {
    discard_mode_ = mode;
  }

  void set_max_size(size_t max_size) {
    max_size_ = max_size;
  }

  void producer_notify_all() {
    que_empty_cv_.notify_all();
  }

private:
  std::deque<T> data_que_;
  mutable std::mutex data_que_mtx_;
  std::condition_variable que_full_cv_, que_empty_cv_;

  std::string name_;
  std::atomic_size_t max_size_ {std::numeric_limits<size_t>::max()};
  DiscardMode discard_mode_ = DiscardMode::POP_FRONT;
};
